using System;

using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Math;

namespace Org.BouncyCastle.Crypto.Engines
{
	/**
	* support class for constructing intergrated encryption ciphers
	* for doing basic message exchanges on top of key agreement ciphers
	*/
	public class IesEngine
	{
		private readonly IBasicAgreement     agree;
		private readonly IDerivationFunction kdf;
		private readonly IMac                mac;
		private readonly BufferedBlockCipher cipher;
		private readonly byte[]              macBuf;

		private bool				forEncryption;
		private ICipherParameters	privParam, pubParam;
		private IesParameters		param;

		/**
		* set up for use with stream mode, where the key derivation function
		* is used to provide a stream of bytes to xor with the message.
		*
		* @param agree the key agreement used as the basis for the encryption
		* @param kdf the key derivation function used for byte generation
		* @param mac the message authentication code generator for the message
		*/
		public IesEngine(
			IBasicAgreement     agree,
			IDerivationFunction kdf,
			IMac                mac)
		{
			this.agree = agree;
			this.kdf = kdf;
			this.mac = mac;
			this.macBuf = new byte[mac.GetMacSize()];
//            this.cipher = null;
		}

		/**
		* set up for use in conjunction with a block cipher to handle the
		* message.
		*
		* @param agree the key agreement used as the basis for the encryption
		* @param kdf the key derivation function used for byte generation
		* @param mac the message authentication code generator for the message
		* @param cipher the cipher to used for encrypting the message
		*/
		public IesEngine(
			IBasicAgreement     agree,
			IDerivationFunction kdf,
			IMac                mac,
			BufferedBlockCipher cipher)
		{
			this.agree = agree;
			this.kdf = kdf;
			this.mac = mac;
			this.macBuf = new byte[mac.GetMacSize()];
			this.cipher = cipher;
		}

		/**
		* Initialise the encryptor.
		*
		* @param forEncryption whether or not this is encryption/decryption.
		* @param privParam our private key parameters
		* @param pubParam the recipient's/sender's public key parameters
		* @param param encoding and derivation parameters.
		*/
		public void Init(
			bool                     forEncryption,
			ICipherParameters            privParameters,
			ICipherParameters            pubParameters,
			ICipherParameters            iesParameters)
		{
			this.forEncryption = forEncryption;
			this.privParam = privParameters;
			this.pubParam = pubParameters;
			this.param = (IesParameters)iesParameters;
		}

		private byte[] DecryptBlock(
			byte[]  in_enc,
			int     inOff,
			int     inLen,
			byte[]  z)
		{
			byte[]          M = null;
			KeyParameter    macKey = null;
			KdfParameters   kParam = new KdfParameters(z, param.GetDerivationV());
			int             macKeySize = param.MacKeySize;

			kdf.Init(kParam);

			inLen -= mac.GetMacSize();

			if (cipher == null)     // stream mode
			{
				byte[] Buffer = GenerateKdfBytes(kParam, inLen + (macKeySize / 8));

				M = new byte[inLen];

				for (int i = 0; i != inLen; i++)
				{
					M[i] = (byte)(in_enc[inOff + i] ^ Buffer[i]);
				}

				macKey = new KeyParameter(Buffer, inLen, (macKeySize / 8));
			}
			else
			{
				int cipherKeySize = ((IesWithCipherParameters)param).CipherKeySize;
				byte[] Buffer = GenerateKdfBytes(kParam, (cipherKeySize / 8) + (macKeySize / 8));

				cipher.Init(false, new KeyParameter(Buffer, 0, (cipherKeySize / 8)));

				M = cipher.DoFinal(in_enc, inOff, inLen);

				macKey = new KeyParameter(Buffer, (cipherKeySize / 8), (macKeySize / 8));
			}

			byte[] macIV = param.GetEncodingV();

			mac.Init(macKey);
			mac.BlockUpdate(in_enc, inOff, inLen);
			mac.BlockUpdate(macIV, 0, macIV.Length);
			mac.DoFinal(macBuf, 0);

			inOff += inLen;

			for (int t = 0; t < macBuf.Length; t++)
			{
				if (macBuf[t] != in_enc[inOff + t])
				{
					throw (new InvalidCipherTextException("IMac codes failed to equal."));
				}
			}

			return M;
		}

		private byte[] EncryptBlock(
			byte[]  input,
			int     inOff,
			int     inLen,
			byte[]  z)
		{
			byte[]          C = null;
			KeyParameter    macKey = null;
			KdfParameters   kParam = new KdfParameters(z, param.GetDerivationV());
			int             c_text_length = 0;
			int             macKeySize = param.MacKeySize;

			if (cipher == null)     // stream mode
			{
				byte[] Buffer = GenerateKdfBytes(kParam, inLen + (macKeySize / 8));

				C = new byte[inLen + mac.GetMacSize()];
				c_text_length = inLen;

				for (int i = 0; i != inLen; i++)
				{
					C[i] = (byte)(input[inOff + i] ^ Buffer[i]);
				}

				macKey = new KeyParameter(Buffer, inLen, (macKeySize / 8));
			}
			else
			{
				int cipherKeySize = ((IesWithCipherParameters)param).CipherKeySize;
				byte[] Buffer = GenerateKdfBytes(kParam, (cipherKeySize / 8) + (macKeySize / 8));

				cipher.Init(true, new KeyParameter(Buffer, 0, (cipherKeySize / 8)));

				c_text_length = cipher.GetOutputSize(inLen);
				byte[] tmp = new byte[c_text_length];

				int len = cipher.ProcessBytes(input, inOff, inLen, tmp, 0);
				len += cipher.DoFinal(tmp, len);

				C = new byte[len + mac.GetMacSize()];
				c_text_length = len;

				Array.Copy(tmp, 0, C, 0, len);

				macKey = new KeyParameter(Buffer, (cipherKeySize / 8), (macKeySize / 8));
			}

			byte[] macIV = param.GetEncodingV();

			mac.Init(macKey);
			mac.BlockUpdate(C, 0, c_text_length);
			mac.BlockUpdate(macIV, 0, macIV.Length);
			//
			// return the message and it's MAC
			//
			mac.DoFinal(C, c_text_length);
			return C;
		}

		private byte[] GenerateKdfBytes(
			KdfParameters	kParam,
			int				length)
		{
			byte[] buf = new byte[length];

			kdf.Init(kParam);

			kdf.GenerateBytes(buf, 0, buf.Length);

			return buf;
		}

		public byte[] ProcessBlock(
			byte[]  input,
			int     inOff,
			int     inLen)
		{
			agree.Init(privParam);

			BigInteger z = agree.CalculateAgreement(pubParam);

			// TODO Check that this is right (...Unsigned? Check length?)
			byte[] zBytes = z.ToByteArray();

			return forEncryption
				?	EncryptBlock(input, inOff, inLen, zBytes)
				:	DecryptBlock(input, inOff, inLen, zBytes);
		}
	}

}
